#include "wake_word_detect.h"
#include "application.h"

#include <esp_log.h>
#include <model_path.h>
#include <arpa/inet.h>
#include <sstream>

// 定义检测运行状态的事件位
#define DETECTION_RUNNING_EVENT 1

// 日志标签
static const char* TAG = "WakeWordDetect";

/**
 * @brief 构造函数
 * 初始化唤醒词检测器并创建事件组
 */
WakeWordDetect::WakeWordDetect()
    : afe_data_(nullptr),
      wake_word_pcm_(),
      wake_word_opus_() {

    event_group_ = xEventGroupCreate();
}

/**
 * @brief 析构函数
 * 释放资源并删除事件组
 */
WakeWordDetect::~WakeWordDetect() {
    if (afe_data_ != nullptr) {
        afe_iface_->destroy(afe_data_);
    }

    if (wake_word_encode_task_stack_ != nullptr) {
        heap_caps_free(wake_word_encode_task_stack_);
    }

    vEventGroupDelete(event_group_);
}

/**
 * @brief 初始化唤醒词检测器
 * @param codec 音频编解码器指针
 */
void WakeWordDetect::Initialize(AudioCodec* codec) {
    codec_ = codec;
    // 确定参考通道数量
    int ref_num = codec_->input_reference() ? 1 : 0;

    // 初始化语音识别模型
    srmodel_list_t *models = esp_srmodel_init("model");
    for (int i = 0; i < models->num; i++) {
        ESP_LOGI(TAG, "Model %d: %s", i, models->model_name[i]);
        // 查找唤醒词模型
        if (strstr(models->model_name[i], ESP_WN_PREFIX) != NULL) {
            wakenet_model_ = models->model_name[i];
            auto words = esp_srmodel_get_wake_words(models, wakenet_model_);
            // 按分号分割获取所有唤醒词
            std::stringstream ss(words);
            std::string word;
            while (std::getline(ss, word, ';')) {
                wake_words_.push_back(word);
            }
        }
    }

    // 构建输入格式字符串，M表示麦克风通道，R表示参考通道
    std::string input_format;
    for (int i = 0; i < codec_->input_channels() - ref_num; i++) {
        input_format.push_back('M');
    }
    for (int i = 0; i < ref_num; i++) {
        input_format.push_back('R');
    }
    // 配置音频前端处理参数
    afe_config_t* afe_config = afe_config_init(input_format.c_str(), models, AFE_TYPE_SR, AFE_MODE_HIGH_PERF);
    afe_config->aec_init = codec_->input_reference();
    afe_config->aec_mode = AEC_MODE_SR_HIGH_PERF;
    afe_config->afe_perferred_core = 1;
    afe_config->afe_perferred_priority = 1;
    afe_config->memory_alloc_mode = AFE_MEMORY_ALLOC_MORE_PSRAM;
    
    // 创建音频前端处理接口和数据
    afe_iface_ = esp_afe_handle_from_config(afe_config);
    afe_data_ = afe_iface_->create_from_config(afe_config);

    // 创建音频检测任务
    xTaskCreate([](void* arg) {
        auto this_ = (WakeWordDetect*)arg;
        this_->AudioDetectionTask();
        vTaskDelete(NULL);
    }, "audio_detection", 4096, this, 3, nullptr);
}

/**
 * @brief 设置唤醒词检测回调函数
 * @param callback 检测到唤醒词时的回调函数
 */
void WakeWordDetect::OnWakeWordDetected(std::function<void(const std::string& wake_word)> callback) {
    wake_word_detected_callback_ = callback;
}

/**
 * @brief 启动唤醒词检测
 */
void WakeWordDetect::StartDetection() {
    xEventGroupSetBits(event_group_, DETECTION_RUNNING_EVENT);
}

/**
 * @brief 停止唤醒词检测
 */
void WakeWordDetect::StopDetection() {
    xEventGroupClearBits(event_group_, DETECTION_RUNNING_EVENT);
    if (afe_data_ != nullptr) {
        afe_iface_->reset_buffer(afe_data_);
    }
}

/**
 * @brief 检查唤醒词检测是否正在运行
 * @return 如果检测正在运行返回true，否则返回false
 */
bool WakeWordDetect::IsDetectionRunning() {
    return xEventGroupGetBits(event_group_) & DETECTION_RUNNING_EVENT;
}

/**
 * @brief 输入音频数据进行唤醒词检测
 * @param data 待检测的音频数据
 */
void WakeWordDetect::Feed(const std::vector<int16_t>& data) {
    afe_iface_->feed(afe_data_, data.data());
}

/**
 * @brief 获取每次Feed所需的数据大小
 * @return 每次Feed所需的采样点数
 */
size_t WakeWordDetect::GetFeedSize() {
    return afe_iface_->get_feed_chunksize(afe_data_) * codec_->input_channels();
}

/**
 * @brief 音频检测任务
 * 持续检测音频流中的唤醒词
 */
void WakeWordDetect::AudioDetectionTask() {
    auto fetch_size = afe_iface_->get_fetch_chunksize(afe_data_);
    auto feed_size = afe_iface_->get_feed_chunksize(afe_data_);
    ESP_LOGI(TAG, "Audio detection task started, feed size: %d fetch size: %d",
        feed_size, fetch_size);

    while (true) {
        // 等待检测运行事件
        xEventGroupWaitBits(event_group_, DETECTION_RUNNING_EVENT, pdFALSE, pdTRUE, portMAX_DELAY);

        // 获取处理后的音频数据
        auto res = afe_iface_->fetch_with_delay(afe_data_, portMAX_DELAY);
        if (res == nullptr || res->ret_value == ESP_FAIL) {
            continue;;
        }

        // 存储唤醒词数据用于语音识别，如识别说话人
        StoreWakeWordData((uint16_t*)res->data, res->data_size / sizeof(uint16_t));

        // 检测到唤醒词
        if (res->wakeup_state == WAKENET_DETECTED) {
            StopDetection();
            last_detected_wake_word_ = wake_words_[res->wake_word_index - 1];

            // 触发唤醒词检测回调
            if (wake_word_detected_callback_) {
                wake_word_detected_callback_(last_detected_wake_word_);
            }
        }
    }
}

/**
 * @brief 存储唤醒词音频数据
 * @param data 音频数据指针
 * @param samples 采样点数
 */
void WakeWordDetect::StoreWakeWordData(uint16_t* data, size_t samples) {
    // 将音频数据存储到wake_word_pcm_
    wake_word_pcm_.emplace_back(std::vector<int16_t>(data, data + samples));
    // 保留约2秒的数据，检测持续时间为32ms（采样率16000，块大小512）
    while (wake_word_pcm_.size() > 2000 / 32) {
        wake_word_pcm_.pop_front();
    }
}

/**
 * @brief 编码唤醒词音频数据
 * 将PCM数据编码为Opus格式
 */
void WakeWordDetect::EncodeWakeWordData() {
    wake_word_opus_.clear();
    if (wake_word_encode_task_stack_ == nullptr) {
        wake_word_encode_task_stack_ = (StackType_t*)heap_caps_malloc(4096 * 8, MALLOC_CAP_SPIRAM);
    }
    // 创建编码任务
    wake_word_encode_task_ = xTaskCreateStatic([](void* arg) {
        auto this_ = (WakeWordDetect*)arg;
        {
            auto start_time = esp_timer_get_time();
            // 创建Opus编码器
            auto encoder = std::make_unique<OpusEncoderWrapper>(16000, 1, OPUS_FRAME_DURATION_MS);
            encoder->SetComplexity(0); // 0是最快的复杂度

            // 编码所有PCM数据
            for (auto& pcm: this_->wake_word_pcm_) {
                encoder->Encode(std::move(pcm), [this_](std::vector<uint8_t>&& opus) {
                    std::lock_guard<std::mutex> lock(this_->wake_word_mutex_);
                    this_->wake_word_opus_.emplace_back(std::move(opus));
                    this_->wake_word_cv_.notify_all();
                });
            }
            this_->wake_word_pcm_.clear();

            auto end_time = esp_timer_get_time();
            ESP_LOGI(TAG, "Encode wake word opus %zu packets in %lld ms",
                this_->wake_word_opus_.size(), (end_time - start_time) / 1000);

            // 添加一个空包表示编码结束
            std::lock_guard<std::mutex> lock(this_->wake_word_mutex_);
            this_->wake_word_opus_.push_back(std::vector<uint8_t>());
            this_->wake_word_cv_.notify_all();
        }
        vTaskDelete(NULL);
    }, "encode_detect_packets", 4096 * 8, this, 2, wake_word_encode_task_stack_, &wake_word_encode_task_buffer_);
}

/**
 * @brief 获取编码后的唤醒词Opus数据
 * @param opus 用于存储获取的Opus数据的向量
 * @return 如果获取到非空数据返回true，否则返回false
 */
bool WakeWordDetect::GetWakeWordOpus(std::vector<uint8_t>& opus) {
    std::unique_lock<std::mutex> lock(wake_word_mutex_);
    // 等待直到有数据可用
    wake_word_cv_.wait(lock, [this]() {
        return !wake_word_opus_.empty();
    });
    opus.swap(wake_word_opus_.front());
    wake_word_opus_.pop_front();
    return !opus.empty();
}
